(* ::Package:: *)

(* ::Text:: *)
(*Description: *)
(*	Utility functions for drawing Feynman diagrams and constructing amplitudes from QGRAF outputs. *)
(**)
(*Author: Zhewen Mo (mozhewen@outlook.com, mozw@itp.ac.cn)*)
(**)
(*Mathematica version: 14.0*)
(**)
(*Last update: 2025.3.25*)


(* ::Section:: *)
(*Begin*)


BeginPackage["FeynGRAF`"]


ClearAll[ImportQGRAF]

ClearAll[$QGRAFExecutable, RunQGRAF]

Unprotect[in, out]
ClearAll[propagatorStyleDict, in, out]
ClearAll[DrawGraph, GraphTopo, FindTopoIsomorphism]

ClearAll[Propagator, Vertex]
ClearAll[ConstructAmp]


Begin["`Private`"]


(* Package directory *)
pd = DirectoryName[$InputFileName]


If[$OperatingSystem === "MacOSX",
	SetEnvironment["PATH" -> Environment["PATH"] <> ":/opt/homebrew/bin"]
]


DeletePath[path_String] :=
	Switch[FileType[path],
		Directory, DeleteDirectory[path, DeleteContents -> True],
		File, DeleteFile[path],
		None, Null
	]


(* ::Section:: *)
(*Import*)


ImportQGRAF::usage =
"ImportQGRAF[file] imports the QGRAF output with the style defined by 'style.yaml'. \
'yq' command (by Mike Farah) is needed. "

ImportQGRAF[file_String] := ImportString[
	RunProcess[{"yq", "-o=json", ".", "-I=0", file}, "StandardOutput"],
	"JSON"
] //. {rs__Rule} :> Association[rs] // Dataset


(* ::Section:: *)
(*Run QGRAF*)


$QGRAFExecutable = FileNameJoin[{$HomeDirectory, "Packages", "qgraf-4.0.5", "qgraf"}]


controlTemp = StringTemplate[
"config = `display-mode`;
output = '`output`';
style = '`style`';

model = '`model`';
in = `in`;
out = `out`;
loops = `loops`;
loop_momentum = `loop_momentum`;
options = \\
    `options`;

`epilog`"]


RunQGRAF::usage =
"RunQGRAF[modelFile, {field1[p1], ...} -> {field'1[k1], ...}, {l, il}] automatically \
generates QGRAF control files and runs QGRAF to get il-loop Feynman diagrams with \
loop momenta labeled by 'l'. "

Options[RunQGRAF] = {
	"TempDir" -> ".temp_FeynGRAF",
	"options" -> {"onshell", "notadpole"},
	Epilog -> "",
	"DisplayMode" -> None
}
RunQGRAF[modelFile_String,
	in_List -> out_List, {l_Symbol, il_Integer},
	OptionsPattern[]] := Module[{
	TempDir = OptionValue["TempDir"],
	DisplayMode = OptionValue["DisplayMode"],
	rt
},
	DeletePath[TempDir];

	CreateDirectory[TempDir];
	FileTemplateApply[controlTemp, <|
		"display-mode" -> Switch[DisplayMode, _String, DisplayMode, _, "noinfo"],
		"output" -> "grafs.yml",
		"style" -> BreakString@FileNameJoin[{pd, "style.yaml"}],
		"model" -> BreakString@ExpandFileName[modelFile],
		"in" -> StringRiffle[TextString/@in, ", "],
		"out" -> StringRiffle[TextString/@out, ", "],
		"loops" -> il,
		"loop_momentum" -> SymbolName[l],
		"options" -> StringRiffle[OptionValue["options"], ", \\\n    "],
		"epilog" -> OptionValue[Epilog]
	|>,
		FileNameJoin[{TempDir, "qgraf.dat"}]
	];
	RunProcess[{$QGRAFExecutable, "qgraf.dat"}, "StandardOutput",
		ProcessDirectory -> TempDir
	] // If[FileExistsQ@FileNameJoin[{TempDir, "qgraf.dat"}] && Head[DisplayMode] =!= String,
		Null,
		Print@StandardForm[#]
	]&;

	ImportQGRAF[FileNameJoin[{TempDir, "grafs.yml"}]]
]

BreakString[str_String] :=
	StringRiffle[StringPartition[str, UpTo[80]], "' \\\n'"]


(* ::Section:: *)
(*Graphs*)


propagatorStyleDict::usage =
"propagatorStyleDict = {
	{\"field\", \"dual-field\"} \[Rule] styleFunc1,
	{in | out, \"field\"} \[Rule] styleFunc2,
	...
}"

propagatorStyleDict = {}

Protect[in, out]


DrawGraph::usage =
"DrawGraph[diagram] draws the Feynman diagrams according to the dataset \
generated by QGRAF with the style file 'style.yaml'. The drawing style of \
the propagators is defined by 'propagatorStyleDict'. 
GraphLayout -> \"SpringElectricalEmbedding\" | ...
\"ExternalLabels\" -> {\"field\" -> func}
"

Options[DrawGraph] = {
	GraphLayout -> "SpringElectricalEmbedding",
	"ExternalLabels" -> {}
}
DrawGraph[list_List, op:OptionsPattern[]] := DrawGraph[#, op]& /@ list
DrawGraph[diagram_Dataset, op:OptionsPattern[]] := With[{diagramN = Normal[diagram]},
	Which[
	Head[diagramN] === List,
		DrawGraph[Dataset[#], op]& /@ diagramN,
	KeyExistsQ[diagram, "diagrams"],
		DrawGraph[diagram["diagrams"], op],
	True,
		Join[
			(
				Labeled[(Replace[{#["field"], #["dual-field"]}, Append[propagatorStyleDict, {_, _} -> Identity]])[
					#["from"] \[DirectedEdge] #["to"]
				], #["momentum"]]
			)& /@ Replace[Normal@diagram["propagators"], Null -> {}],
			(
				Labeled[(Replace[{ToExpression[#["type"]], #["field"]}, Append[propagatorStyleDict, {_, _} -> Identity]])[
					Which[
						#["type"] === "in", #["field_index"] \[DirectedEdge] #["to"],
						#["type"] === "out", #["from"] \[DirectedEdge] #["field_index"]
					]
				], #["momentum"]]
			)& /@ Replace[Normal@diagram["legs"], Null -> {}]
		] // EdgeTaggedGraph[#,
			VertexLabels -> (With[{lbl = Replace[#["field"],
					Append[OptionValue["ExternalLabels"], _ -> (Null&)]
				][#["field_index"]]
			},
				If[lbl === Null, Nothing, #["field_index"] -> lbl]
			]& /@ Replace[Normal@diagram["legs"], Null -> {}]),
			GraphLayout -> OptionValue[GraphLayout]
		]&
	]
]


GraphTopo::usage =
"GraphTopo[diagram] only generates the topological information of the Feynman diagrams for checking \
graph isomorphism. 
NOTE: The keys in \"Vertices\" are guaranteed to be sorted. "

GraphTopo[list_List] := GraphTopo /@ list
GraphTopo[diagram_Dataset] := With[{diagramN = Normal[diagram]},
	Which[
	Head[diagramN] === List,
		GraphTopo[Dataset[#]]& /@ diagramN,
	KeyExistsQ[diagram, "diagrams"],
		GraphTopo[diagram["diagrams"]],
	True,
		Association@{
			"Vertices" -> KeySort@GroupBy[
				{#["fields"], #["index"]}& /@ Replace[Normal@diagram["vertices"], Null -> {}],
				First -> Last
			],
			"Edges" -> Join[
				If[#["field"] === #["dual-field"],
					UndirectedEdge[#["from"], #["to"], #["field"]],
				(*Else*)
					DirectedEdge[#["from"], #["to"], {#["field"], #["dual-field"]}]
				]& /@ Replace[Normal@diagram["propagators"], Null -> {}],
				Which[
					#["type"] === "in", DirectedEdge[#["field_index"], #["to"], {in, #["field"]}],
					#["type"] === "out", DirectedEdge[#["from"], #["field_index"], {out, #["field"]}]
				]& /@ Replace[Normal@diagram["legs"], Null -> {}]
			]
		}
	]
]


ue[a_, b_, t_] /; Not@OrderedQ[{a, b}] := ue[b, a, t]
de[a_, b_, t_]

FindTopoIsomorphism[extList_List:{}][
	<|"Vertices" -> v1_, "Edges" -> e1_|>,
	<|"Vertices" -> v2_, "Edges" -> e2_|>] := Module[{idx1, perms2},
	Catch[
		If[EchoLabel["type1: "][Length/@v1] === EchoLabel["type2: "][Length/@v2] (* Vertex types match *),
			idx1 = Join[extList, Catenate@Values[v1]];
			perms2 = Catenate/@Tuples[
				Prepend[Values[Permutations/@v2], Permutations[extList]]
			];
			Do[With[{rl = Echo@Thread[idx1 -> idx2]},
				If[Total[e1/.rl/.{UndirectedEdge -> ue, DirectedEdge -> de}] ===
					Total[e2/.{UndirectedEdge -> ue, DirectedEdge -> de}],
					Echo["Succeed"]; Throw[rl]
				]
			], {idx2, perms2}];
		];
		Null
	]
]


(* ::Section:: *)
(*Amplitudes*)


ConstructAmp[list_List] := ConstructAmp /@ list
ConstructAmp[diagram_Dataset] :=  With[{diagramN = Normal[diagram]},
	Which[Head[diagramN] === List,
		ConstructAmp[Dataset[#]]& /@ diagramN,
	KeyExistsQ[diagram, "diagrams"],
		ConstructAmp[diagram["diagrams"]],
	True,
		Times[
			ToExpression[diagram["signsym"]],
			Times@@((Vertex[
				#["fields"],
				#["field_indices"],
				ToExpression[#["momenta"]]
			])& /@ Replace[Normal@diagram["vertices"], Null -> {}]),
			Times@@((Propagator[
				{#["field"], #["dual-field"]},
				{#["field_index"], #["dual-field_index"]},
				ToExpression[#["momentum"]]
			])& /@ Replace[Normal@diagram["propagators"], Null -> {}])
		]
	]
]


(* ::Section:: *)
(*End*)


End[]
EndPackage[]
